from typing import Optional, List
from fastapi import HTTPException
from tenacity import retry, stop_after_attempt, wait_exponential, wait_fixed, retry_if_exception_type
from weaviate.classes.query import Filter
from weaviate.classes.config import Property, DataType, Configure
# from weaviate.collections.classes.config import Configure, Property, DataType
from pydantic import BaseModel
from db.my_weaviate import DB
from utils.main import get_embedding

class VectorTable:
    @staticmethod
    def get_chunks(collection_name: str, source: str, offset: int, limit: int,keyword: str = None):
        """
        根据来源获取文档分块信息
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        # 构建基础过滤条件 by ysh 2025-1-9 10:25
        filters = Filter.by_property("source").equal(source)

        # 如果提供了关键词，添加模糊检索条件 by ysh 2025-1-9 10:25
        if keyword:
            keyword_filter = Filter.by_property("text").like(f"*{keyword}*")
            filters = filters & keyword_filter  # 组合过滤条件

         # 获取符合条件的总数
        total_count_response = collection.query.fetch_objects(
            offset=0,
            filters=filters
        )
        total_count = len(total_count_response.objects)

        # 获取分页后的对象
        response = collection.query.fetch_objects(
            limit=limit,
            offset=offset,
            filters=filters
        )

        result = {
            "total_count": total_count,  # 获取总数
            "objects": response.objects  # 分页后的对象
        }
        if response is None:
            raise HTTPException(status_code=404, detail="文件不存在")
        return {"code": 200, "message": "文件分块信息获取成功","total_count": total_count, "data": response.objects}

    @staticmethod
    async def update_chunk(collection_name: str, chunk_id: str, text: str, description: str, embed: str):
        """
        更新文档分块信息
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        # 重新生成向量
        text_vector = await get_embedding(embed, text)
        description_vector = await get_embedding(embed, description)

        collection.data.update(
            uuid=chunk_id,
            properties={
                "text": text,
                "description": description
            },
            vector={
                "text": text_vector[0],
                "description": description_vector[0]
            }
        )

    @staticmethod
    def stop_chunk(collection_name: str, chunk_id: str):
        """
        停止文档分块信息
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        collection.data.update(
            uuid=chunk_id,
            properties={
                "stop": True
            }
        )

    @staticmethod
    def delete_chunk_by_id(collection_name: str, chunk_id: str):
        """
        删除文档分块信息
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        collection.data.delete_by_id(
            uuid=chunk_id
        )

    @staticmethod
    async def delete_chunks(collection_name: str, source: str | int):
        """
        删除文档分块信息
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        collection.data.delete_many(
            where=Filter.by_property("source").equal(source)
        )

    @staticmethod
    @retry(stop=stop_after_attempt(3),
           wait=wait_fixed(2),
           retry=retry_if_exception_type(Exception))
    async def create_chunk(
            collection_name: str,
            filename: str,
            text: str,
            description: str,
            embed: str,
            num: int = 0
    ):
        """
        创建文本块
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        # 生成向量
        text_vector = await get_embedding(model=embed, text=text)
        description_vector = await get_embedding(model=embed, text=description)

        collection.data.insert(
            properties={
                "source": filename,
                "text": text,
                "description": description,
                "num": num,
                "stop": False
            },
            vector={
                "text": text_vector[0],
                "description": description_vector[0]
            }
        )

    @staticmethod
    async def search_chunks_by_vector(
            collection_name: str,
            query_vector: list,
            limit: int = 40
    ) -> list:
        """
        根据向量搜索文本块
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        response = collection.query.near_vector(
            filters=Filter.by_property("stop").equal(False),
            near_vector=query_vector,
            target_vector=["text", "description"],
            limit=limit,
        )

        return response.objects

    @staticmethod
    async def search_chunks_by_bm25(
            collection_name: str,
            query_keywords: str,
            limit: int = 40
    ) -> list:
        """
        根据关键词进行bm25检索
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        response = collection.query.bm25(
            query=query_keywords,
            query_properties=["text", "description"],
            filters=Filter.by_property("stop").equal(False),
            limit=limit,
        )

        return response.objects

    @staticmethod
    async def search_chunks_by_hybrid(
            collection_name: str,
            query_vector: list,
            query_keywords: str,
            alpha: float,
            limit: int = 40
    ) -> list:
        """
        根据向量和关键词进行混合检索
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        response = collection.query.hybrid(
            query=query_keywords,
            vector=query_vector,
            query_properties=["text", "description"],
            target_vector=["text", "description"],
            alpha=alpha,
            filters=Filter.by_property("stop").equal(False),
            limit=limit,
        )

        return response.objects

    @staticmethod
    def get_chunk_by_num(collection_name: str, source: str, num: int) -> list:
        """
        根据num获取文本块
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        response = collection.query.fetch_objects(
            filters=Filter.by_property("stop").equal(False)
                    & Filter.by_property("source").equal(source)
                    & Filter.by_property("num").equal(num),
        )

        return response.objects

    @staticmethod
    def delete_collection(collection_name: str):
        """
        删除集合
        """
        client = DB.get_client()
        client.collections.delete(collection_name)

    @staticmethod
    def create_collection(collection_name: str, vectorizer_config=None, properties=None, is_base: bool = False):
        """
        创建集合
        """
        client = DB.get_client()
        if is_base:
            client.collections.create(
                name=collection_name,
                vectorizer_config=[
                    Configure.NamedVectors.none(name="text"),
                    Configure.NamedVectors.none(name="description"),
                ],
                properties=[
                    Property(name="source", data_type=DataType.TEXT),
                    Property(name="text", data_type=DataType.TEXT),
                    Property(name="description", data_type=DataType.TEXT),
                    Property(name="num", data_type=DataType.INT),
                    Property(name="stop", data_type=DataType.BOOL),
                ]
            )
        else:
            client.collections.create(name=collection_name, properties=properties, vectorizer_config=vectorizer_config)

    @staticmethod
    async def create_chunk_generic(
            collection_name: str,
            properties: dict,
            vector: dict | list,
    ):
        """
        创建文本块
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)

        uuid = collection.data.insert(
            properties=properties,
            vector=vector,
        )
        return uuid

    @staticmethod
    async def update_chunk_generic(
            collection_name: str,
            chunk_id: str,
            properties: dict,
            vector: dict | list,
    ):
        """
        更新文档分块信息
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)

        collection.data.update(
            uuid=chunk_id,
            properties=properties,
            vector=vector,
        )

    @staticmethod
    async def get_chunk_by_filter(collection_name: str, filters: Filter):
        """
        根据筛选条件分块信息
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        response = collection.query.fetch_objects(
            filters=filters
        )

        return response.objects

    @staticmethod
    async def search_chunks_by_hybrid_generic(
            collection_name: str,
            query_vector: list,
            query_keywords: str,
            query_properties: list,
            target_vector: list,
            alpha: float,
            limit: int = 40
    ) -> list:
        """
        根据向量和关键词进行混合检索
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        response = collection.query.hybrid(
            query=query_keywords,
            query_vector=query_vector,
            query_properties=query_properties,
            target_vector=target_vector,
            alpha=alpha,
            filters=Filter.by_property("stop").equal(False),
            limit=limit,
        )

        return response.objects

    @staticmethod
    def exists(collection_name: str) -> bool:
        """
        检查向量库是否已存在
        """
        client = DB.get_client()
        collection = client.collections.get(collection_name)
        return collection.exists()


Vectors = VectorTable()
